# import mindspore as ms
import numpy as np
import torch
from einops import rearrange
import mindspore as ms
from mindspore import load_checkpoint, save_checkpoint, context
import os
from omegaconf import OmegaConf
from models import build_model, build_backbone
from models_torch import build_model_torch, convert_ms_checkpoint


context.set_context(mode=context.PYNATIVE_MODE, device_target="CPU",)
args = OmegaConf.load('configs/default.yaml')

ms_model = build_backbone(args)
save_checkpoint(ms_model , 'test.ckpt') # save without decoder
ms_ckpt = load_checkpoint('test.ckpt')
os.remove('test.ckpt')
converted_dict = convert_ms_checkpoint(ms_ckpt)
torch_model = build_model_torch(args)
msg = torch_model.load_state_dict(converted_dict, strict=False)
print(msg.missing_keys)

##### testing consistency
# ms_model.set_train(False)
# torch_model.eval()

# manager = ms_model.manager
# modality = manager[0]
# batch_size = 4
# lo, hi = modality.get_range()
# c, p, g = modality.num_channels, modality.patch_size, modality.grid_size
# image = np.random.rand(batch_size, modality.num_channels, *modality.img_size)
# image_patchify = rearrange(
#         image, 
#         'b c (g0 p0) (g1 p1) (g2 p2) -> b (g0 g1 g2) (c p0 p1 p2)', 
#         g0=g[0], g1=g[1], g2=g[2], p0=p[0], p1=p[1], p2=p[2]
#         )
# inputs = np.zeros((batch_size, manager.seq_len, manager.total_dim))
# inputs[:, :modality.seq_len, lo:hi] = image_patchify
# instance_ids = np.zeros((batch_size, manager.seq_len + 1))
# instance_ids[:, :modality.seq_len + 1] =  1

# m_idx_ms = ms.Tensor(np.zeros(batch_size), dtype=ms.int32)
# inputs_ms = ms.Tensor(inputs, dtype=ms.float32)
# instance_ids_ms = ms.Tensor(instance_ids, dtype=ms.int32)
# ms_outputs = ms_model(inputs_ms, m_idx_ms, instance_ids_ms)
# ms_outputs = ms_outputs.asnumpy()
# ms_outputs = ms_outputs[:, :modality.seq_len + 1]
# print(ms_outputs.shape)

# with torch.no_grad():
#     inputs_torch = torch.from_numpy(image).float()
#     outputs_torch = torch_model(inputs_torch, 0)
#     outputs_torch = outputs_torch.numpy()
# print(outputs_torch.shape)

# max_ms, max_torch = np.max(np.abs(ms_outputs)), np.max(np.abs(outputs_torch))
# print(f'Max MS: {max_ms}, Max Torch: {max_torch}')
# l0 = np.max(np.abs(ms_outputs - outputs_torch))
# l1 = np.sum(np.abs(ms_outputs - outputs_torch))
# print(f'Max Diff: {l0}, Sum Abs Diff: {l1}')
# rel = l0 / max(max_ms, max_torch)
# print(f'Rel Diff: {rel}')
